from math import log
from math import sqrt
from math import exp
import numpy as np

def poly(x):
	return 2*x[0] -3*x[1]**2

def objPoly(x):
	valueToReach = 5
	return abs(poly(x)-valueToReach)


def cma_es(objFunction, x0, s0, stopFitness, nbMaxIteration = None):
	N = len(x0)
	x0 = np.array(x0)
	xmean = x0
	s0 = np.array(s0)
	sigma = s0
	if nbMaxIteration == None:
		nbMaxIteration = 1000*N**2

	#Parametre concernant la selection
	nbTirage = 4 + 3*int(round(log(N))) #lambda
	mu = nbTirage/2
	weights = log(mu + 0.5)
	weights = np.array([x+weights for x in map(log, np.arange(1, mu+1))])
	weights = weights/weights.sum()
	mueff = weights.sum()**2 / sum(weights**2)

	#Parametre concernant l'adaptation
	cc = (4 + mueff/N) / (N+4 + 2*mueff/N)
	cs = (mueff+2) / (N+mueff+5)
	c1 = 2 / ((N+1.3)**2 + mueff)
	cmu = 2 * (mueff-2+1/mueff) / ((N+2)**2+mueff)
	damps = 1 + 2*max(0, sqrt((mueff-1)/(N+1))-1)

	#Parametres internes
	pc = np.zeros(N)
	ps = np.zeros(N)
	B = np.matrix(np.eye(N,N))
	D = np.ones(N)
	C = B * np.matrix(np.diag(D**2)) * B.transpose()
	invsqrtC = B * np.matrix(np.diag(D**-1)) * B.transpose()
	eigeneval = 0
	chiN=N**0.5*(1-1/(4*N)+1/(21*N**2))

	#Debut iterations
	counteval = 0;
	arx = np.matrix(np.zeros((N, nbTirage)))
	arfitness = np.zeros(nbTirage)
	while counteval < nbMaxIteration:
		for k in xrange(nbTirage):


			#print "sigma :", sigma
			#print "B :", B
			#print "D :", D
			#print "np.random.normal(0, 1, N) :", np.random.normal(0, 1, N)
			#print "D * np.random.normal(0, 1, N) :", D *np.random.normal(0, 1, N)
			#print "B * D... :",B*np.matrix( D *np.random.normal(0, 1, N)).transpose()
			#print "sigma * B * D... :",sigma * B*np.matrix( D *np.random.normal(0, 1, N)).transpose()
			a = np.matrix(x0).transpose() + sigma * B * np.matrix(D * np.random.normal(0, 1, N)).transpose()
			arx[:,k] = np.matrix(x0).transpose() + sigma * B*np.matrix( D *np.random.normal(0, 1, N)).transpose()
			arfitness[k] = objFunction(arx[:,k])
			counteval += 1
		#arfitness.sort() #Il faut faire un sort dans l'autre sens
		xold = np.matrix(xmean).reshape((N,1))
		mu = int(round(mu))
		arxSorted = arx[:, arfitness.argsort()]
		xmean = arxSorted[:,0:mu] * np.matrix(weights).transpose()
		xmean = xmean.reshape((N,1))
		print "xmean :", xmean, xmean.shape
		print "xold :", xold, xold.shape
		#Evolution path
		ps = (1-cs)*ps + sqrt(cs*(2-cs)*mueff) * invsqrtC * (xmean-xold) / sigma;
		#print "ps :", ps
		hsig = np.linalg.norm(ps)/sqrt(1-(1-cs)**(2*counteval/nbTirage))/chiN < 1.4 + 2/(N+1)
		#print "hsig :", ps
		pc = (1-cc)*pc + hsig * sqrt(cc*(2-cc)*mueff) * (xmean-xold) / sigma;
		print "xmean-xold :", xmean-xold
		#Adaptation de la matrice de covariance
		artmp = (1.0/sigma) * (arxSorted[:,0:mu] - np.array([xold for i in xrange(mu)]).transpose())
		print "artimp ok"
		C = (1-c1-cmu) * C + c1 * (pc*pc.transpose()+ (1-hsig) * cc*(2-cc) * C) + cmu * np.matrix(artmp) * np.matrix(np.diag(weights)) * np.matrix(artmp).transpose()
		print "=================="
		print 'C :', C
		sigma = sigma * exp((cs/damps)*(np.linalg.norm(ps)/chiN - 1));
		print "=================="
		print 'sigma :', sigma
		if counteval - eigeneval > nbTirage/(c1+cmu)/N/10:
			eigeneval = counteval
			C = np.triu(C) + np.triu(C,1).transpose()
			D = np.linalg.eig(C);
			B = D[1] #vecteurs propres norme
			D = D[0] #valeurs propres
			D = np.sqrt(np.diag(D))
			print np.matrix(B) * np.matrix(np.diag(np.array(D)**-1)).transpose()
			#invsqrtC = np.matrix(B).reshape((1,N)) * np.matrix(np.diag(np.array(D)**-1)) * np.matrix(B).transpose() #TODO je m'arrete la 
			#print "invsqrtC :", invsqrtC



if __name__ == "__main__":
	x0 = [3,5]
	print "Solution pour", x0, ":", poly(x0)
	s0 = 2
	cma_es(objPoly, x0, s0, 0.01)
